from typing import Dict, Any, List, Optional

from langchain.vectorstores.base import VectorStore
from langchain_community.vectorstores import FAISS, Chroma, Pinecone
from langchain_core.retrievers import BaseRetriever


class ModularRAG:
    """模块化RAG系统"""

    def __init__(
            self,
            name: str,
            vectorstore: VectorStore,
            retriever: BaseRetriever,
            config: Dict[str, Any] = None
    ):
        self.name = name
        self.vectorstore = vectorstore
        self.retriever = retriever
        self.config = config or {}

    def retrieve(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """检索相关文档"""
        try:
            docs = self.retriever.get_relevant_documents(query)
            return [{"content": doc.page_content, "metadata": doc.metadata} for doc in docs[:k]]
        except Exception as e:
            return [{"error": str(e)}]

    def add_documents(self, documents: List[Dict[str, Any]]):
        """添加文档到向量库"""
        try:
            # 这里需要根据具体的向量库实现
            pass
        except Exception as e:
            print(f"添加文档失败: {e}")


class RAGFactory:
    """RAG工厂类"""

    _vectorstore_types = {
        "faiss": FAISS,
        "chroma": Chroma,
        "pinecone": Pinecone
    }

    @classmethod
    def create_rag(
            cls,
            rag_type: str,
            embedding_model,
            config: Dict[str, Any] = None
    ) -> Optional[ModularRAG]:
        """创建RAG实例"""
        try:
            vectorstore_class = cls._vectorstore_types.get(rag_type, FAISS)

            # 创建向量库
            vectorstore = vectorstore_class.from_texts(
                texts=[""],  # 初始化空向量库
                embedding=embedding_model,
                **(config or {})
            )

            # 创建检索器
            retriever = vectorstore.as_retriever()

            return ModularRAG(
                name=f"{rag_type}_rag",
                vectorstore=vectorstore,
                retriever=retriever,
                config=config
            )
        except Exception as e:
            print(f"创建RAG失败: {e}")
            return None

    @classmethod
    def create_rags_from_list(
            cls,
            rag_list: List[Dict[str, Any]],
            embedding_model
    ) -> List[ModularRAG]:
        """从RAG列表创建RAG实例"""
        rags = []
        for rag_config in rag_list:
            rag_type = rag_config.get("type", "faiss")
            config = rag_config.get("config", {})
            rag = cls.create_rag(rag_type, embedding_model, config)
            if rag:
                rags.append(rag)
        return rags

    @classmethod
    def get_available_types(cls) -> Dict[str, str]:
        """获取可用的RAG类型"""
        return {
            "faiss": "FAISS向量库",
            "chroma": "Chroma向量库",
            "pinecone": "Pinecone向量库"
        }

    @classmethod
    def register_vectorstore(cls, store_type: str, store_class: type):
        """注册新的向量库类型"""
        cls._vectorstore_types[store_type] = store_class
